#############################################################################################
#############################################################################################
#
#   Open Source License/Disclaimer, Forecast Systems Laboratory NOAA/OAR/GSD,
#   325 Broadway Boulder, CO 80305
#
#   This software is distributed under the Open Source Definition, which may be
#   found at http://www.opensource.org/.
#   In particular, redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following conditions are met:
#
#   - Redistributions of source code must retain this notice, this list of
#     conditions and the following disclaimer.
#
#   - Redistributions in binary form must provide access to this notice, this
#     list of conditions and the following disclaimer, and the underlying source code.
#
#   - All modifications to this software must be clearly documented, and are
#     solely the responsibility of the agent making the modifications.
#
#   - If significant modifications or enhancements are made to this software,
#     the GSD Software Policy Manager (softwaremgr.fsl@noaa.gov) should be notified.
#
#   THIS SOFTWARE AND ITS DOCUMENTATION ARE IN THE PUBLIC DOMAIN AND ARE
#   FURNISHED "AS IS." THE AUTHORS, THE UNITED STATES GOVERNMENT, ITS INSTRUMENTALITIES,
#   OFFICERS, EMPLOYEES, AND AGENTS MAKE NO WARRANTY, EXPRESS OR IMPLIED, AS TO
#   THE USEFULNESS OF THE SOFTWARE AND DOCUMENTATION FOR ANY PURPOSE. THEY ASSUME
#   NO RESPONSIBILITY (1) FOR THE USE OF THE SOFTWARE AND DOCUMENTATION; OR (2) TO PROVIDE
#   TECHNICAL SUPPORT TO USERS.
#############################################################################################
#############################################################################################

###############################################################################
###############################################################################
# This library contains a suite of functions to generate simulations
# of Lorenz 63 model, also included are the tangent linear and adjoint functions.
###############################################################################
###############################################################################

import numpy as np

# Import generic Model Class from Shallow Water Model
from ..Model.Model_Class import ModelClass

class Lorenz63ModelClass(ModelClass):
    
    def __init__(self,start_time,stop_time,Ntimestep,sigma=10.,beta=8./3,rho=28.0, \
                 ic_start_range=-15.,ic_end_range=15.,x_noise=1,y_noise=1, \
                 z_noise=1,seed=None,time_diff_choice='leapfrog'):
        
        ModelClass.__init__(self,start_time,stop_time,Ntimestep,time_diff_choice=time_diff_choice)
        
        self.Name = 'Lorenz 63 Model'
        self.VarDim = 3  # We only have 3 coordinates (x,y,z) in three-dimensional space
        
        #*******Parameter Block
        self.sigma=sigma	        # the system parameters
        self.beta = beta	        # the system parameters
        self.rho  = rho                 # the system parameters
        
        # model configuration
        self.ic_start_range=ic_start_range   # The location range of IC ---left
        self.ic_end_range=ic_end_range       # The location range of IC ---right
        self.seed = seed
        
        self.x_noise=x_noise
        self.y_noise=y_noise
        self.z_noise=z_noise
       
        # Define IC randomly
        np.random.seed(self.seed)
        ic =  self.ic_start_range + (self.ic_end_range - self.ic_start_range) * np.random.random((1,3))
        
        # We need to convert them to 2-dim arraies
        self.x = np.array([[ic[0][0]]])
        self.y = np.array([[ic[0][1]]])
        self.z = np.array([[ic[0][2]]])
        
        # Stack all variables into one array
        self.ic = np.vstack((self.x,self.y,self.z))
        
        # Define the noise diagnoal matrix
        self.noise = np.eye((self.x.size+self.y.size+self.z.size))
        
        for i in range(self.x.size):
            self.noise[i][i] = self.x_noise
        for i in range(self.x.size,self.x.size+self.y.size):
            self.noise[i][i] = self.y_noise
        for i in range(self.x.size+self.y.size,self.x.size+self.y.size+self.z.size):
            self.noise[i][i] = self.z_noise
            
    # Self description
    def __repr__(self):
        return '*'*50+'\n'+ \
                self.Name +'\n'+ \
                'IC: x,y,z= '+str(np.vsplit(self.ic,self.VarDim))+'\n' + \
                'Using '+self.time_diff_choice+' time differential scheme to propagate\n'+ \
                '*'*50
   
    # Initilize pertubation
    def model_init_tl(self):
        self.g_x = np.zeros_like(self.x)
        self.g_y = np.zeros_like(self.y)
        self.g_z = np.zeros_like(self.z)
        
        # Stack all variables into one array
        self.ic_tl = np.vstack((self.g_x,self.g_y,self.g_z))
    
    # Initilize gradient
    def model_init_ad(self):
        self.a_x = np.zeros_like(self.x)
        self.a_y = np.zeros_like(self.y)
        self.a_z = np.zeros_like(self.z)
        
        # Stack all variables into one array
        self.ic_ad = np.vstack((self.a_x,self.a_y,self.a_z))
        
    # Lorenz 63 model -- Add normal noise on IC
    def model_ic_add_noise(self):
        
        self.g_x = self.x + self.x_noise * np.random.randn(self.x.shape[0],self.x.shape[1])
        self.g_y = self.y + self.y_noise * np.random.randn(self.y.shape[0],self.x.shape[1])
        self.g_z = self.z + self.z_noise * np.random.randn(self.z.shape[0],self.x.shape[1])

        # Stack all variables into one array
        self.ic_perturbation = np.vstack((self.g_x,self.g_y,self.g_z))
                

    # Lorenz 63 Model Tendency
    def model_tendency(self, ic, time) :
    
        [self.x,self.y,self.z] = np.vsplit(ic, self.VarDim)
        
        dxdt = self.sigma * (self.y - self.x)
        dydt = self.x * (self.rho - self.z) - self.y
        dzdt = self.x * self.y - self.beta * self.z
        
        return np.vstack((dxdt,dydt,dzdt))

    
    # Lorenz 63 Model Tendency --- Tangent Linear
    def model_tendency_tl(self, ic, ic_tl, time) :
          
        [self.x,self.y,self.z] = np.vsplit(ic, self.VarDim)
        [self.g_x,self.g_y,self.g_z] = np.vsplit(ic_tl, self.VarDim)
        
        dxdt = self.sigma * (self.y - self.x)
        dydt = self.x * (self.rho - self.z) - self.y
        dzdt = self.x * self.y - self.beta * self.z
        
        dg_xdt = self.sigma * (self.g_y - self.g_x)
        dg_ydt = self.g_x * (self.rho - self.z) - self.x * self.g_z - self.g_y
        dg_zdt = self.g_x * self.y + self.x * self.g_y - self.beta * self.g_z
        
        return np.vstack((dxdt,dydt,dzdt)), np.vstack((dg_xdt,dg_ydt,dg_zdt))
    
    
    # Lorenz 63 Model Tendency --- Adjoint
    def model_tendency_ad(self, basic_states, forcing, time) :
    
        da_xdt = np.zeros_like(self.a_x)
        da_ydt = np.zeros_like(self.a_y)
        da_zdt = np.zeros_like(self.a_z)  
        
        self.a_x.fill(0.)
        self.a_y.fill(0.)
        self.a_z.fill(0.)
        
        [self.x, self.y, self.z] = np.vsplit(basic_states, self.VarDim)
        
        [forcing_x, forcing_y, forcing_z] = np.vsplit(forcing, self.VarDim)
        da_xdt = da_xdt + forcing_x
        da_ydt = da_ydt + forcing_y
        da_zdt = da_zdt + forcing_z
        
        self.a_x = (self.rho - self.z)*da_ydt -self.sigma*da_xdt + self.y*da_zdt
        self.a_y = self.sigma*da_xdt - da_ydt + self.x*da_zdt
        self.a_z = -(self.x*da_ydt) -self.beta*da_zdt
 			
        da_xdt.fill(0.)
        da_ydt.fill(0.)   
        da_zdt.fill(0.)
        
        return np.vstack((self.a_x,self.a_y,self.a_z))